import pandas as pd
from decision_tree_id3 import DecisionTree

# 创建示例数据集
data = {
    'Outlook': ['Sunny', 'Sunny', 'Overcast', 'Rain', 'Rain', 'Rain', 'Overcast'],
    'Temperature': ['Hot', 'Hot', 'Hot', 'Mild', 'Cool', 'Cool', 'Cool'],
    'Humidity': ['High', 'High', 'High', 'High', 'Normal', 'Normal', 'Normal'],
    'Wind': ['Weak', 'Strong', 'Weak', 'Weak', 'Weak', 'Strong', 'Strong'],
    'PlayTennis': ['No', 'No', 'Yes', 'Yes', 'Yes', 'No', 'Yes']
}

df = pd.DataFrame(data)

# 分离特征和目标
X = df.drop('PlayTennis', axis=1)
y = df['PlayTennis']

# 创建并训练C4.5决策树
c45_tree = DecisionTree(algorithm='c4.5')
c45_tree.fit(X, y)

# 打印决策树规则
print("决策树规则:")
c45_tree.print_tree()

# 创建新数据并预测
new_data = pd.DataFrame({
    'Outlook': ['Rain'],
    'Temperature': ['Mild'],
    'Humidity': ['Normal'],
    'Wind': ['Strong']
})

prediction = c45_tree.predict(new_data)
print(f"\n预测结果: {prediction[0]}")

# 可视化决策树 (需要安装graphviz)
graph = c45_tree.to_graphviz()
graph.render('tennis_decision_tree', format='png', cleanup=True)
print("决策树已保存为 tennis_decision_tree.png")